{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "workbookDir: /home/adrisuau/qlm_notebook/notebook/qaths/notebooks\n"
     ]
    }
   ],
   "source": [
    "try:\n",
    "    import qaths\n",
    "except ImportError:\n",
    "    # If we can't find qaths, then try to add some directory where it could be.\n",
    "    # First, we determine the current notebook directory.\n",
    "    import os\n",
    "    if not 'workbookDir' in globals():\n",
    "        workbookDir = os.getcwd()\n",
    "    \n",
    "    def is_root(path: str) -> bool:\n",
    "        path = os.path.realpath(path)\n",
    "        return path == os.path.dirname(path)\n",
    "    \n",
    "    # Then we try to find qaths in a parent directory.\n",
    "    import sys\n",
    "    current_dir = os.path.dirname(os.path.abspath(workbookDir))\n",
    "    import_successfull = False\n",
    "    # Loop in parent directories until we can import qaths or we \n",
    "    # find the root.\n",
    "    while not (import_successfull or is_root(current_dir)):\n",
    "        sys.path.append(current_dir)\n",
    "        try:\n",
    "            import qaths\n",
    "        except ImportError:\n",
    "            # Remove the added directory from the PYTHONPATH in order to\n",
    "            # not pollute it with a lot of useless directories.\n",
    "            sys.path.pop()\n",
    "            current_dir = os.path.dirname(current_dir)\n",
    "        else:\n",
    "            print(\"Found qaths library in {}.\".format(current_dir))\n",
    "            import_successfull = True\n",
    "    \n",
    "    # If qaths has not been successfully imported, warn the user.\n",
    "    if not import_successfull:\n",
    "        print(\"qaths not found!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qaths.applications.wave_equation.utils import construct_Hamiltonians_1D, compute_qubit_number_from_considered_points_1D\n",
    "from qaths.generation.integer_weighted import matrix2int_weighted_permutation\n",
    "from qaths.qram import QRAM\n",
    "\n",
    "from qaths.applications.wave_equation.oracles import (\n",
    "    permutation_oracle_1,\n",
    "    permutation_oracle_2,\n",
    "    weight_oracle_1,\n",
    "    weight_oracle_2,\n",
    "    sign_oracle_1,\n",
    "    sign_oracle_2,\n",
    ")\n",
    "\n",
    "\n",
    "discretisation_points_number = 6\n",
    "\n",
    "Hs = construct_Hamiltonians_1D(discretisation_points_number)\n",
    "\n",
    "considered_points = discretisation_points_number - 2\n",
    "n = compute_qubit_number_from_considered_points_1D(considered_points)\n",
    "\n",
    "# Hs entries have only 0 and 1 as **weights** (sign is handled by another qubit) and\n",
    "# so 1 qubit is enough to encode the weights\n",
    "int_size = 1\n",
    "oracles = [matrix2int_weighted_permutation(H, int_size) for H in Hs]\n",
    "\n",
    "permutations = (oracle[0] for oracle in oracles)\n",
    "weights = (oracle[1] for oracle in oracles)\n",
    "signs = (oracle[2] for oracle in oracles)\n",
    "\n",
    "Ms1 = [QRAM(permutation) for permutation in permutations]\n",
    "Vs1 = [QRAM(weight) for weight in weights]\n",
    "Ss1 = [QRAM(sign) for sign in signs]\n",
    "\n",
    "Ms2 = [\n",
    "    permutation_oracle_1(n, discretisation_points_number),\n",
    "    permutation_oracle_2(n, discretisation_points_number),\n",
    "]\n",
    "Vs2 = [\n",
    "    weight_oracle_1(n, discretisation_points_number),\n",
    "    weight_oracle_2(n, discretisation_points_number),\n",
    "]\n",
    "Ss2 = [\n",
    "    sign_oracle_1(n, discretisation_points_number),\n",
    "    sign_oracle_2(n, discretisation_points_number),\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from qat.lang.AQASM import Program, H\n",
    "from qat.core.task import Task\n",
    "#Running on linalg.\n",
    "from qat.linalg import get_qpu_server\n",
    "\n",
    "def all_outputs(M, V, S):\n",
    "    \n",
    "    ma_size = M.arity - 2*n\n",
    "    va_size = V.arity - n - int_size\n",
    "    sa_size = S.arity - n - 1\n",
    "    \n",
    "    prog = Program()\n",
    "    x = prog.qalloc(n)\n",
    "    m = prog.qalloc(n)\n",
    "    w = prog.qalloc(int_size)\n",
    "    s = prog.qalloc(1)\n",
    "    ancillas = prog.qalloc(max(ma_size, va_size, sa_size))\n",
    "    \n",
    "    print(ma_size, va_size, sa_size, n)\n",
    "    \n",
    "    for i in range(n):\n",
    "        prog.apply(H, x[i])\n",
    "    \n",
    "    prog.apply(M, x, m, ancillas[:ma_size])\n",
    "    prog.apply(V, x, w, ancillas[:va_size])\n",
    "    prog.apply(S, x, s, ancillas[:sa_size])\n",
    "    \n",
    "    circ = prog.to_circ()\n",
    "    \n",
    "    task = Task(circ, get_qpu_server())\n",
    "    task.execute()\n",
    "\n",
    "    found_states = set()\n",
    "    for state in task.states():\n",
    "        if state.probability > 1e-10:\n",
    "            str_state = str(state.state)[1:-1]\n",
    "            found_states.add(\"{x} {m} {v} {s} {a}\".format(x=str_state[:n][::-1],\n",
    "                                                          m=str_state[n:2*n][::-1], \n",
    "                                                          v=str_state[2*n], \n",
    "                                                          s=str_state[2*n+1], \n",
    "                                                          a=str_state[2*n+2:][::-1]))\n",
    "    return found_states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 0 0 4\n",
      "4 5 0 4\n",
      "|0000 0101 1 0 >\n",
      "|0001 0110 1 0 >\n",
      "|0010 0111 1 0 >\n",
      "|0011 1000 1 0 >\n",
      "|0100 0100 0 0 >\n",
      "|0101 0000 1 0 >\n",
      "|0110 0001 1 0 >\n",
      "|0111 0010 1 0 >\n",
      "|1000 0011 1 0 >\n",
      "|1001 1001 0 0 >\n",
      "|1010 1010 0 0 >\n",
      "|1011 1011 0 0 >\n",
      "|1100 1100 0 0 >\n",
      "|1101 1101 0 0 >\n",
      "|1110 1110 0 0 >\n",
      "|1111 1111 0 0 >\n",
      "================================================================================\n",
      "|0000 0101 1 0 00000>\n",
      "|0001 0110 1 0 00000>\n",
      "|0010 0111 1 0 00000>\n",
      "|0011 1000 1 0 00000>\n",
      "|0100 1001 0 0 00000>\n",
      "|0101 0000 1 0 00000>\n",
      "|0110 0001 1 0 00000>\n",
      "|0111 0010 1 0 00000>\n",
      "|1000 0011 1 0 00000>\n",
      "|1001 0100 0 0 00000>\n",
      "|1010 1010 0 0 00000>\n",
      "|1011 1011 0 0 00000>\n",
      "|1100 1100 0 0 00000>\n",
      "|1101 1101 0 0 00000>\n",
      "|1110 1110 0 0 00000>\n",
      "|1111 1111 0 0 00000>\n"
     ]
    }
   ],
   "source": [
    "found1 = all_outputs(Ms1[1], Vs1[1], Ss1[1])\n",
    "found2 = all_outputs(Ms2[1], Vs2[1], Ss2[1])\n",
    "\n",
    "found1 = sorted(list(found1))\n",
    "found2 = sorted(list(found2))\n",
    "\n",
    "for found in found1:\n",
    "    print(\"|{}>\".format(found))\n",
    "print(\"=\"*80)\n",
    "for found in found2:\n",
    "    print(\"|{}>\".format(found))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.4.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
